{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "collapsed": true,
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "from itertools import *\n",
    "from functools import *\n",
    "import operator\n",
    "import joblib\n",
    "import datetime\n",
    "from scipy.stats import norm\n",
    "from multiprocessing import Manager"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "outputs": [],
   "source": [
    "# 最大值函数\n",
    "def I(x: np.ndarray):\n",
    "    t = np.zeros([x.shape[0], 2])\n",
    "    t[:, 0] = x\n",
    "    return np.max(t, axis=1)\n",
    "\n",
    "# 期望函数\n",
    "def E(x: np.ndarray, p: np.ndarray):\n",
    "    return np.dot(x, p)\n",
    "\n",
    "# 依照正态分布计算概率值\n",
    "def p_norm(x: np.ndarray, mu, sigma):\n",
    "    init_arr = norm.cdf(x, mu, sigma)\n",
    "    res_arr = np.zeros(x.shape[0])\n",
    "    res_arr[0] = init_arr[0]\n",
    "    # res_arr[0] = init_arr[0] - norm.cdf(x[0] - 1, mu, sigma)\n",
    "    res_arr[1:] = np.diff(init_arr)\n",
    "    res_arr[-1] = 1 - init_arr[-2]\n",
    "    return res_arr\n",
    "\n",
    "# Variable\n",
    "alpha = 0.95\n",
    "# Kt\n",
    "mu_K_t = 6  # 均值\n",
    "cv_K_t = 0.3  # 变异系数\n",
    "sigma_K_t = cv_K_t * mu_K_t  # 标准差\n",
    "K_t = np.array([4, 5, 6, 7, 8])\n",
    "p_K_t = p_norm(K_t, mu_K_t, sigma_K_t)\n",
    "# Dt\n",
    "mu_t = 5\n",
    "sigma_t = 0.5\n",
    "D_t = np.array(list(range(0, 10)))\n",
    "p_D_t = p_norm(D_t, mu_t, sigma_t)\n",
    "p = 19\n",
    "b = 19\n",
    "c_e = 6\n",
    "c_m = 2\n",
    "h_e = 1\n",
    "h_m = 0.2\n",
    "# theta = h_m / (1 - alpha)\n",
    "theta = 4"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "outputs": [],
   "source": [
    "# 多线程优化模块 (有线程池)\n",
    "from concurrent.futures import ThreadPoolExecutor, as_completed\n",
    "\n",
    "_max_worker = 5\n",
    "\n",
    "def grid_search_exa(iter_grid, func, params: tuple) -> list:\n",
    "    pool = ThreadPoolExecutor(max_workers=_max_worker)\n",
    "    tasks = [pool.submit(func, *params, x, y) for x, y in iter_grid]\n",
    "    results = [t.result() for t in as_completed(tasks)]\n",
    "    return results\n",
    "\n",
    "\n",
    "def linear_search_exa(iter_arr, func, params: tuple) -> list:\n",
    "    pool = ThreadPoolExecutor(max_workers=_max_worker)\n",
    "    tasks = [pool.submit(func, *params, x) for x in iter_arr]\n",
    "    results = [t.result() for t in as_completed(tasks)]\n",
    "    return results"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "outputs": [],
   "source": [
    "# 基础算子\n",
    "def psi_1T(v_1, v_2, v_3, e, z):\n",
    "    return h_e * E(I(v_1 - e - D_t), p_D_t) \\\n",
    "           + b * E(I(D_t - v_1 + e), p_D_t) \\\n",
    "           + c_e * (v_1 - e) \\\n",
    "           + c_m * (z - v_3) \\\n",
    "           + h_m * (v_3 - v_1) \\\n",
    "           + theta * e\n",
    "\n",
    "def psi_2T(v_2, v_3, y,z):\n",
    "    return h_e * E(I(y - D_t), p_D_t) \\\n",
    "           + b * E(I(D_t - y), p_D_t) \\\n",
    "           + c_e * y \\\n",
    "           + c_m * (z - v_3) \\\n",
    "           + h_m * (v_3 - y)\n",
    "\n",
    "def psi_3T(v_3, y, z):\n",
    "    return h_e * E(I(y - D_t), p_D_t) \\\n",
    "           + b * E(I(D_t - y), p_D_t) \\\n",
    "           + c_e * y \\\n",
    "           + c_m * (z - v_3) \\\n",
    "           + h_m * (v_3 - y)"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "outputs": [],
   "source": [
    "# psi_1t\n",
    "def psi_1t(v_1, v_2, v_3, e, z):\n",
    "    # 先计算基础值\n",
    "    val = psi_1T(v_1, v_2, v_3, e, z)\n",
    "    # 再从缓存寻找上一阶段的 f_t+1\n",
    "    last_cache = global_cache[epoch]\n",
    "    _D_t = [last_cache[f\"psi_1t_f_{v_1 - e - D_t_i}_{v_2 - e - D_t_i}_{v_3 - e - D_t_i}_{z - e - D_t_i}\"][\"value\"] for D_t_i in D_t]\n",
    "    return val + alpha * E(_D_t, p_D_t)\n",
    "# psi_2t\n",
    "def psi_2t(v_2, v_3, y, z):\n",
    "    val = psi_2T(v_2, v_3, y, z)\n",
    "    last_cache = global_cache[epoch]\n",
    "    _D_t = [last_cache[f\"psi_2t_f_{y - D_t_i}_{v_2 - D_t_i}_{v_3 - D_t_i}_{z - D_t_i}\"][\"value\"] for D_t_i in D_t]\n",
    "    return val + alpha * E(_D_t, p_D_t)\n",
    "# psi_3t\n",
    "def psi_3t(v_3, y, z):\n",
    "    val = psi_3T(v_3, y, z)\n",
    "    last_cache = global_cache[epoch]\n",
    "    _D_t = [last_cache[f\"psi_3t_f_{y - D_t_i}_{y - D_t_i}_{v_3 - D_t_i}_{z - D_t_i}\"][\"value\"] for D_t_i in D_t]\n",
    "    return val + alpha * E(_D_t, p_D_t)\n",
    "# psi_1t(1, 2, 3, 0, 4)\n",
    "# psi_2t(3, 4, 2, 4)\n",
    "# psi_3t(4, 3, 4)"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "outputs": [],
   "source": [
    "def phi_1t(v_0, v_1, v_2, v_3, K_t_i):\n",
    "    e_range = range(0, v_1 - v_0 +1)\n",
    "    z_range = range(v_3, v_3 + K_t_i)\n",
    "    e_z_grid = product(e_range, z_range)\n",
    "    psi_1t_values = grid_search_exa(e_z_grid, psi_1t, (v_1, v_2, v_3))\n",
    "    return - c_e * v_0 + np.min(psi_1t_values)\n",
    "\n",
    "def phi_2t(v_0, v_1, v_2, v_3, K_t_i):\n",
    "    y_range = range(v_1, v_2 + 1)\n",
    "    z_range = range(v_3, v_3 + K_t_i)\n",
    "    y_z_grid = product(y_range, z_range)\n",
    "    psi_2t_values = grid_search_exa(y_z_grid, psi_2t, (v_2, v_3))\n",
    "    return - c_e * v_0 + np.min(psi_2t_values)\n",
    "\n",
    "def phi_3t(v_0, v_1, v_2, v_3, K_t_i):\n",
    "    y_range = range(v_2, v_3 + 1)\n",
    "    z_range = range(v_3, v_3 + K_t_i)\n",
    "    y_z_grid = product(y_range, z_range)\n",
    "    psi_3t_values = grid_search_exa(y_z_grid, psi_3t, (v_3, ))\n",
    "    return - c_e * v_0 + np.min(psi_3t_values)\n",
    "\n",
    "def f_t(v_0, v_1, v_2, v_3):\n",
    "    _K_t = []\n",
    "    for K_t_i in K_t:\n",
    "        exp_1_value = phi_1t(v_0, v_1, v_2, v_3, K_t_i)\n",
    "        exp_2_value = phi_2t(v_0, v_1, v_2, v_3, K_t_i)\n",
    "        exp_3_value = phi_3t(v_0, v_1, v_2, v_3, K_t_i)\n",
    "        _K_t.append(np.min([exp_1_value, exp_2_value, exp_3_value]))\n",
    "    return E(_K_t, p_K_t)"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "outputs": [],
   "source": [
    "# Hint: 生成 psi_1t 在 T = 1, 2, 3 阶段所需要计算的全部 f_t\n",
    "def psi_1t_init(v_0, v_1, v_2, v_3, epoch):\n",
    "    psi_f_dict = {}\n",
    "    for K_t_i in K_t:\n",
    "        e = range(0, v_1 - v_0 + 1)\n",
    "        z = range(v_3, v_3 + K_t_i + 1)\n",
    "        for D_t_i in D_t:\n",
    "            for _e, _z in product(e, z):\n",
    "                if epoch == max_epoch:\n",
    "                    cache_state = True\n",
    "                else:\n",
    "                    cache_state = False\n",
    "                key = f\"psi_1t_f_{v_1 - _e - D_t_i}_{v_2 - _e - D_t_i}_{v_3 - _e - D_t_i}_{_z - _e - D_t_i}\"\n",
    "                psi_f_dict[key] = {\n",
    "                    \"value\": 0,\n",
    "                    \"params\": (v_1 - _e - D_t_i, v_2 - _e - D_t_i, v_3 - _e - D_t_i, _z - _e - D_t_i),\n",
    "                    \"cached\": cache_state\n",
    "                }\n",
    "    return psi_f_dict\n",
    "\n",
    "\n",
    "def psi_2t_init(v_0, v_1, v_2, v_3, epoch):\n",
    "    psi_f_dict = {}\n",
    "    for K_t_i in K_t:\n",
    "        y_12 = range(v_1, v_2 + 1)\n",
    "        z = range(v_3, v_3 + K_t_i + 1)\n",
    "        for D_t_i in D_t:\n",
    "            for _y_12, _z in product(y_12, z):\n",
    "                if epoch == max_epoch:\n",
    "                    cache_state = True\n",
    "                else:\n",
    "                    cache_state = False\n",
    "                key = f\"psi_2t_f_{_y_12 - D_t_i}_{v_2 - D_t_i}_{v_3 - D_t_i}_{_z - D_t_i}\"\n",
    "                psi_f_dict[key] = {\n",
    "                    \"value\": 0,\n",
    "                    \"params\": (_y_12 - D_t_i, v_2 - D_t_i, v_3 - D_t_i, _z - D_t_i),\n",
    "                    \"cached\": cache_state\n",
    "                }\n",
    "    return psi_f_dict\n",
    "\n",
    "\n",
    "def psi_3t_init(v_0, v_1, v_2, v_3, epoch):\n",
    "    psi_f_dict = {}\n",
    "    for K_t_i in K_t:\n",
    "        y_23 = range(v_2, v_3 + 1)\n",
    "        z = range(v_3, v_3 + K_t_i + 1)\n",
    "        for D_t_i in D_t:\n",
    "            for _y_23, _z in product(y_23, z):\n",
    "                if epoch == max_epoch:\n",
    "                    cache_state = True\n",
    "                else:\n",
    "                    cache_state = False\n",
    "                key = f\"psi_3t_f_{_y_23 - D_t_i}_{_y_23 - D_t_i}_{v_3 - D_t_i}_{_z - D_t_i}\"\n",
    "                psi_f_dict[key] = {\n",
    "                    \"value\": 0,\n",
    "                    \"params\": (_y_23 - D_t_i, _y_23 - D_t_i, v_3 - D_t_i, _z - D_t_i),\n",
    "                    \"cached\": cache_state,\n",
    "                }\n",
    "    return psi_f_dict"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "模型初始化完成. 耗时  6.65 秒\n"
     ]
    }
   ],
   "source": [
    "import time\n",
    "\n",
    "init_begin_time = time.perf_counter()\n",
    "\n",
    "global_cache = {}\n",
    "epoch = 1\n",
    "max_epoch = 2\n",
    "v_0 = 1\n",
    "v_1 = 4\n",
    "v_2 = 6\n",
    "v_3 = 8\n",
    "# 在开始迭代前, 先给初始条件的期望子式赋 0\n",
    "# Hint: 初值 v0, v1, v2, v3 均已确定, 因而可以生成全部的 e, z, y 取值范围\n",
    "# 由于要推导的是到 epoch = 1 的 f_t, 所以需要将 epoch = 2, 3, ... 等的对应表达式初始化为 0\n",
    "\n",
    "global_cache[epoch] = {}\n",
    "global_cache[epoch].update(psi_1t_init(v_0, v_1, v_2, v_3, epoch))\n",
    "global_cache[epoch].update(psi_2t_init(v_0, v_1, v_2, v_3, epoch))\n",
    "global_cache[epoch].update(psi_3t_init(v_0, v_1, v_2, v_3, epoch))\n",
    "\n",
    "while True:\n",
    "    if epoch == max_epoch:\n",
    "        break\n",
    "\n",
    "    global_cache[epoch + 1] = {}\n",
    "    # 遍历上一轮所有需要计算的 f_t\n",
    "    for key in global_cache[epoch].keys():\n",
    "        params = global_cache[epoch][key][\"params\"]\n",
    "        global_cache[epoch + 1].update(psi_1t_init(*params, epoch))\n",
    "        global_cache[epoch + 1].update(psi_2t_init(*params, epoch))\n",
    "        global_cache[epoch + 1].update(psi_3t_init(*params, epoch))\n",
    "    epoch = epoch + 1\n",
    "\n",
    "init_end_time = time.perf_counter()\n",
    "print(f\"模型初始化完成. 耗时 {init_end_time - init_begin_time: .2f} 秒\")"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "71.31434036663416\n",
      "计算完成, 耗时  43.72 秒\n"
     ]
    }
   ],
   "source": [
    "# 在模型初始化之后, 需要逐个填充对应的缓存\n",
    "calc_begin_time = time.perf_counter()\n",
    "\n",
    "epoch = max_epoch\n",
    "while True:\n",
    "    if epoch == 1:\n",
    "        print(f_t(v_0, v_1, v_2, v_3))\n",
    "        break\n",
    "\n",
    "    def func(key):\n",
    "        params = global_cache[epoch - 1][key][\"params\"]\n",
    "        cached = global_cache[epoch - 1][key][\"cached\"]\n",
    "        return {\"key\": key, \"value\": f_t(*params)}\n",
    "\n",
    "    pool = ThreadPoolExecutor(max_workers=_max_worker)\n",
    "    tasks = [pool.submit(func, key) for key in global_cache[epoch - 1].keys()]\n",
    "    results = [t.result() for t in as_completed(tasks)]\n",
    "    for result in results:\n",
    "        global_cache[epoch - 1][result[\"key\"]][\"value\"] = result[\"value\"]\n",
    "\n",
    "    epoch = epoch - 1\n",
    "\n",
    "calc_end_time = time.perf_counter()\n",
    "print(f\"计算完成, 耗时 {calc_end_time - calc_begin_time: .2f} 秒\")"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "outputs": [
    {
     "data": {
      "text/plain": "['caches/global_cache_T=3_time=2024-01-16 13-08-31.pkl']"
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Epoch = 2, T = 3, Seconds = 57 s\n",
    "# Epoch = 3, T = 4, Seconds = 672 s\n",
    "# Epoch = 4, T = 5, Seconds = 3747.08 s\n",
    "# Epoch = 5, T = 6, Seconds = 60426 s\n",
    "joblib.dump(global_cache, f\"caches/global_cache_T={max_epoch + 1}_time={datetime.datetime.now().strftime('%Y-%m-%d %H-%M-%S')}.pkl\")"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}